# -*- coding: utf-8 -*-
import json
import os
import random

import openai
import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates

app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")


with open(os.path.join(os.path.abspath("."), "gpt_key.txt"), encoding="utf-8") as kf:
    OPENAI_API_KEY = kf.read().strip()


client = openai.OpenAI(
    # This is the default and can be omitted
    api_key=OPENAI_API_KEY,
)


@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    return templates.TemplateResponse("chat.html", {"request": request})


@app.get("/change")
async def change_key(key: str):
    if "sk-" in key:
        with open(
            os.path.join(os.path.abspath("."), "gpt_key.txt"), "w", encoding="utf-8"
        ) as kf:
            kf.write(key)
        global OPENAI_API_KEY
        OPENAI_API_KEY = key
        openai.api_key = key
        return "success"
    else:
        return "err key"


@app.post("/chat")
async def chat(request: Request):
    form = await request.form()
    messages = form["prompts"]
    # api_key = form.get("apiKey")

    # gpt-4
    models = ["gpt-4"]

    if not messages:
        return JSONResponse(
            status_code=400,
            content={
                "error": {
                    "message": "请输入prompts！",
                    "type": "invalid_request_error",
                    "code": "",
                }
            },
        )

    prompts = json.loads(messages)

    try:
        stream = client.chat.completions.create(
            messages=prompts,
            model=random.choice(models),
        )
    except openai.OpenAIError as e:
        return JSONResponse(
            status_code=500,
            content={
                "error": {
                    "message": str(e),
                    "type": "openai_error",
                    "code": e.__class__.__name__,
                }
            },
        )

    def generate():
        for chunk in stream:
            yield chunk.choices[0].delta.content or ""

    return "".join(generate())


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8900)
